from typing import TypedDict
from langgraph.graph import StateGraph,START
from langgraph.types import StreamWriter


class State(TypedDict):
    topic:str
    joke:str

def refine_topic(state:State):
    return {"topic":state["topic"] + " and cats"}

def generate_joke(state:State,writer:StreamWriter):
    writer({"custom_key":"周宇的消息"})
    return {"joke":f"This is a joke about {state['topic']}"}

graph = (StateGraph(State)
         .add_node(refine_topic).add_node(generate_joke)
    .add_edge(START,"refine_topic")
    .add_edge("refine_topic","generate_joke")
    .compile())

for chunk in graph.stream({"topic":"ice_cream"},stream_mode="custom"):
    print(chunk)

